"""
数据生成器API
支持自定义字段生成测试数据并发送到Kafka
"""

from fastapi import APIRouter, HTTPException, BackgroundTasks, Request
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any, Union
import asyncio
import json
import uuid
import random
import time
from datetime import datetime, timedelta
from faker import Faker
import threading
from loguru import logger

from app.core.kafka_client import get_kafka_client
from app.api.websocket import manager as websocket_manager

router = APIRouter(prefix="/api/data-generator", tags=["数据生成器"])

# 全局变量存储当前生成任务
current_task: Optional[Dict] = None
task_lock = threading.Lock()

class FieldConfig(BaseModel):
    """字段配置"""
    name: str = Field(..., description="字段名称")
    type: str = Field(..., description="字段类型")
    options: Optional[Dict[str, Any]] = Field(default={}, description="字段选项配置")

class DataGeneratorConfig(BaseModel):
    """数据生成器配置"""
    topic: str = Field(..., description="目标Topic")
    fields: List[FieldConfig] = Field(..., description="字段配置列表")
    count: int = Field(default=100, ge=1, le=100000, description="生成数据数量")
    interval: int = Field(default=1000, ge=100, le=60000, description="发送间隔(毫秒)")
    batch_size: int = Field(default=1, ge=1, le=1000, description="批次大小")
    stream_mode: bool = Field(default=True, description="是否流式发送")

class GeneratorProgress(BaseModel):
    """生成进度信息"""
    task_id: str
    topic: str
    total_count: int
    sent_count: int
    status: str  # running, completed, failed, stopped
    start_time: datetime
    error_message: Optional[str] = None
    progress_percent: float = 0.0

class DataGenerator:
    """数据生成器"""
    
    def __init__(self):
        self.faker = Faker(['zh_CN', 'en_US'])  # 支持中英文
        
    def generate_field_value(self, field_config: FieldConfig) -> Any:
        """根据字段配置生成值"""
        field_type = field_config.type.lower()
        options = field_config.options or {}
        
        try:
            if field_type == "string":
                min_length = options.get("min_length", 5)
                max_length = options.get("max_length", 20)
                if "values" in options:
                    return random.choice(options["values"])
                return self.faker.text(max_nb_chars=random.randint(min_length, max_length))
                
            elif field_type == "int":
                min_val = options.get("min", 0)
                max_val = options.get("max", 1000)
                return random.randint(min_val, max_val)
                
            elif field_type == "float":
                min_val = options.get("min", 0.0)
                max_val = options.get("max", 1000.0)
                precision = options.get("precision", 2)
                return round(random.uniform(min_val, max_val), precision)
                
            elif field_type == "boolean":
                return random.choice([True, False])
                
            elif field_type == "datetime":
                format_str = options.get("format", "%Y-%m-%d %H:%M:%S")
                if options.get("current", False):
                    return datetime.now().strftime(format_str)
                start_date = options.get("start_date", datetime.now() - timedelta(days=30))
                end_date = options.get("end_date", datetime.now())
                if isinstance(start_date, str):
                    start_date = datetime.fromisoformat(start_date)
                if isinstance(end_date, str):
                    end_date = datetime.fromisoformat(end_date)
                random_date = self.faker.date_time_between(start_date=start_date, end_date=end_date)
                return random_date.strftime(format_str)
                
            elif field_type == "uuid":
                return str(uuid.uuid4())
                
            elif field_type == "email":
                return self.faker.email()
                
            elif field_type == "phone":
                return self.faker.phone_number()
                
            elif field_type == "address":
                return self.faker.address()
                
            elif field_type == "name":
                return self.faker.name()
                
            elif field_type == "company":
                return self.faker.company()
                
            elif field_type == "url":
                return self.faker.url()
                
            elif field_type == "ip":
                return self.faker.ipv4()
                
            elif field_type == "user_agent":
                return self.faker.user_agent()
                
            else:
                return f"unknown_type_{field_type}"
                
        except Exception as e:
            logger.warning(f"生成字段 {field_config.name} 值失败: {str(e)}")
            return f"error_{field_type}"
    
    def generate_record(self, fields: List[FieldConfig]) -> Dict[str, Any]:
        """生成一条记录"""
        record = {}
        for field in fields:
            record[field.name] = self.generate_field_value(field)
        
        # 添加元数据
        record["_generated_at"] = datetime.now().isoformat()
        record["_generator_id"] = "kmsg-ui-generator"
        
        return record

async def send_generated_data(
    task_id: str,
    config: DataGeneratorConfig,
    generator: DataGenerator
):
    """发送生成的数据，通过WebSocket实时推送进度"""
    kafka_client = get_kafka_client()

    try:
        # 初始化任务状态
        with task_lock:
            global current_task
            current_task = {
                "task_id": task_id,
                "topic": config.topic,
                "total_count": config.count,
                "sent_count": 0,
                "status": "running",
                "start_time": datetime.now(),
                "error_message": None
            }

        # 发送初始进度
        await websocket_manager.broadcast_json({
            "type": "data_generator_progress",
            "data": {
                "task_id": task_id,
                "topic": config.topic,
                "total_count": config.count,
                "sent_count": 0,
                "status": "running",
                "progress_percent": 0.0,
                "start_time": current_task["start_time"].isoformat()
            }
        })

        sent_count = 0
        batch_records = []

        for i in range(config.count):
            # 检查任务是否被停止
            with task_lock:
                if current_task is None or current_task.get("status") == "stopped":
                    logger.info(f"任务 {task_id} 被停止")
                    await websocket_manager.broadcast_json({
                        "type": "data_generator_progress",
                        "data": {
                            "task_id": task_id,
                            "status": "stopped",
                            "sent_count": sent_count,
                            "progress_percent": round((sent_count / config.count) * 100, 1)
                        }
                    })
                    return

            # 生成记录
            record = generator.generate_record(config.fields)
            batch_records.append(record)

            # 批量发送或单条发送
            if len(batch_records) >= config.batch_size or i == config.count - 1:
                for record in batch_records:
                    try:
                        # 发送到Kafka
                        await kafka_client.send_message(
                            topic=config.topic,
                            key=record.get("id", str(uuid.uuid4())),
                            value=json.dumps(record, ensure_ascii=False)
                        )
                        sent_count += 1

                        # 更新任务状态
                        with task_lock:
                            if current_task:
                                current_task["sent_count"] = sent_count

                        # 每发送1条消息都推送进度（实时更新）
                        if True:  # 每条消息都更新进度
                            progress_percent = round((sent_count / config.count) * 100, 1)
                            await websocket_manager.broadcast_json({
                                "type": "data_generator_progress",
                                "data": {
                                    "task_id": task_id,
                                    "topic": config.topic,
                                    "total_count": config.count,
                                    "sent_count": sent_count,
                                    "status": "running",
                                    "progress_percent": progress_percent
                                }
                            })

                    except Exception as e:
                        logger.error(f"发送消息失败: {str(e)}")
                        with task_lock:
                            if current_task:
                                current_task["status"] = "failed"
                                current_task["error_message"] = str(e)

                        await websocket_manager.broadcast_json({
                            "type": "data_generator_progress",
                            "data": {
                                "task_id": task_id,
                                "status": "failed",
                                "error_message": str(e),
                                "sent_count": sent_count,
                                "progress_percent": round((sent_count / config.count) * 100, 1)
                            }
                        })
                        return

                batch_records = []

                # 流式发送时的间隔（毫秒转秒）
                if config.stream_mode and i < config.count - 1:
                    await asyncio.sleep(config.interval / 1000.0)

        # 任务完成
        with task_lock:
            if current_task:
                current_task["status"] = "completed"
                current_task["sent_count"] = sent_count

        await websocket_manager.broadcast_json({
            "type": "data_generator_progress",
            "data": {
                "task_id": task_id,
                "topic": config.topic,
                "total_count": config.count,
                "sent_count": sent_count,
                "status": "completed",
                "progress_percent": 100.0
            }
        })

        logger.info(f"数据生成任务 {task_id} 完成，共发送 {sent_count} 条消息")

    except Exception as e:
        logger.error(f"数据生成任务 {task_id} 失败: {str(e)}")
        with task_lock:
            if current_task:
                current_task["status"] = "failed"
                current_task["error_message"] = str(e)

        await websocket_manager.broadcast_json({
            "type": "data_generator_progress",
            "data": {
                "task_id": task_id,
                "status": "failed",
                "error_message": str(e),
                "sent_count": sent_count if 'sent_count' in locals() else 0,
                "progress_percent": round((sent_count / config.count) * 100, 1) if 'sent_count' in locals() else 0
            }
        })
    finally:
        # 清理任务状态
        with task_lock:
            current_task = None



@router.post("/start", response_model=Dict[str, str])
async def start_data_generation(
    config: DataGeneratorConfig,
    background_tasks: BackgroundTasks
):
    """启动数据生成任务"""
    try:
        logger.info(f"收到数据生成请求: {config}")
        logger.info(f"字段配置: {config.fields}")
        logger.info(f"间隔配置: {config.interval} 毫秒")
        # 检查是否已有任务在运行
        with task_lock:
            if current_task and current_task.get("status") == "running":
                raise HTTPException(status_code=400, detail="已有数据生成任务在运行，请先停止当前任务")

        # 验证Topic是否存在
        kafka_client = get_kafka_client()
        topics = await kafka_client.list_topics()
        if config.topic not in topics:
            raise HTTPException(status_code=400, detail=f"Topic '{config.topic}' 不存在")

        # 创建任务ID
        task_id = str(uuid.uuid4())

        # 启动后台任务
        generator = DataGenerator()
        background_tasks.add_task(send_generated_data, task_id, config, generator)

        logger.info(f"启动数据生成任务 {task_id}: Topic={config.topic}, Count={config.count}")

        return {
            "task_id": task_id,
            "message": "数据生成任务已启动",
            "topic": config.topic,
            "count": str(config.count)
        }

    except Exception as e:
        logger.error(f"启动数据生成任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"启动任务失败: {str(e)}")

@router.get("/status")
async def get_current_status():
    """获取当前生成任务状态"""
    with task_lock:
        if current_task is None:
            return {"status": "idle", "message": "无运行中的任务"}

        return {
            "status": "active",
            "task": {
                "task_id": current_task["task_id"],
                "topic": current_task["topic"],
                "total_count": current_task["total_count"],
                "sent_count": current_task["sent_count"],
                "status": current_task["status"],
                "start_time": current_task["start_time"].isoformat(),
                "error_message": current_task["error_message"],
                "progress_percent": round((current_task["sent_count"] / current_task["total_count"]) * 100, 1) if current_task["total_count"] > 0 else 0
            }
        }

@router.post("/stop")
async def stop_current_task():
    """停止当前数据生成任务"""
    with task_lock:
        if current_task is None:
            raise HTTPException(status_code=400, detail="没有运行中的任务")

        if current_task["status"] in ["completed", "failed", "stopped"]:
            raise HTTPException(status_code=400, detail="任务已结束")

        current_task["status"] = "stopped"
        task_id = current_task["task_id"]

    logger.info(f"停止数据生成任务 {task_id}")
    return {"message": "任务已停止"}

@router.post("/debug", response_model=Dict[str, Any])
async def debug_data_generation(request_data: Dict[str, Any]):
    """调试数据生成请求"""
    try:
        logger.info(f"收到调试请求: {request_data}")

        # 尝试验证配置
        config = DataGeneratorConfig(**request_data)
        logger.info(f"验证成功: {config}")

        return {"status": "success", "config": config.dict()}

    except Exception as e:
        logger.error(f"验证失败: {str(e)}")
        return {"status": "error", "error": str(e), "data": request_data}

@router.get("/field-types")
async def get_supported_field_types():
    """获取支持的字段类型"""
    return {
        "field_types": [
            {
                "type": "string",
                "description": "字符串",
                "options": {
                    "min_length": "最小长度",
                    "max_length": "最大长度",
                    "values": "可选值列表"
                }
            },
            {
                "type": "int",
                "description": "整数",
                "options": {
                    "min": "最小值",
                    "max": "最大值"
                }
            },
            {
                "type": "float",
                "description": "浮点数",
                "options": {
                    "min": "最小值",
                    "max": "最大值",
                    "precision": "小数位数"
                }
            },
            {
                "type": "boolean",
                "description": "布尔值",
                "options": {}
            },
            {
                "type": "datetime",
                "description": "日期时间",
                "options": {
                    "format": "日期格式",
                    "current": "使用当前时间",
                    "start_date": "开始日期",
                    "end_date": "结束日期"
                }
            },
            {
                "type": "uuid",
                "description": "UUID",
                "options": {}
            },
            {
                "type": "email",
                "description": "邮箱地址",
                "options": {}
            },
            {
                "type": "phone",
                "description": "电话号码",
                "options": {}
            },
            {
                "type": "address",
                "description": "地址",
                "options": {}
            },
            {
                "type": "name",
                "description": "姓名",
                "options": {}
            },
            {
                "type": "company",
                "description": "公司名称",
                "options": {}
            },
            {
                "type": "url",
                "description": "网址",
                "options": {}
            },
            {
                "type": "ip",
                "description": "IP地址",
                "options": {}
            },
            {
                "type": "user_agent",
                "description": "用户代理",
                "options": {}
            }
        ]
    }
